from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from core.connectors.db_connector import DatabaseConnector

class PostgresqlConnector(DatabaseConnector):
    def __init__(self, host="localhost", user="", password="", database="", port=5432, pool_size=5):


        v = [str(item) for item in [host, user, database, port]]
        un = "_".join(v)
        super().__init__(un)

        """
        初始化 PostgreSQL 数据库连接池
        :param host: 主机地址
        :param user: 用户名
        :param password: 密码
        :param database: 数据库名
        :param port: 端口号
        :param pool_size: 连接池大小
        """
        self.engine = create_engine(
            f"postgresql://{user}:{password}@{host}:{port}/{database}",
            pool_size=pool_size,
            max_overflow=10,
            pool_recycle=1800,
            pool_pre_ping=True,
            echo=False,
        )
        self.Session = sessionmaker(bind=self.engine)

    def get_db_type(self):
        return 'Postgresql'

    def test_connection(self):
        """
        测试数据库连接是否成功。
        :return: True 如果连接成功；False 和错误信息如果连接失败。
        """
        try:
            # 尝试连接并执行简单查询
            with self.engine.connect() as connection:
                connection.execute(text("SELECT 1"))
            return True
        except Exception as e:
            return False, str(e)

    def get_session(self):
        """
        获取 SQLAlchemy 会话
        :return: Session 对象
        """
        return self.Session()

    def get_tables_info(self):
        """
        获取 PostgreSQL 数据库中的表名、记录数和字段个数
        :return: 包含表名、记录数和字段个数的列表
        """
        tables_info = []
        with self.engine.connect() as conn:
            # 获取所有表名
            tables_query = text("""
                SELECT tablename FROM pg_tables WHERE schemaname = 'public';
            """)
            tables = conn.execute(tables_query).fetchall()

            for table_row in tables:
                table = table_row[0]
                # 获取记录数
                row_count_query = text(f"SELECT COUNT(*) FROM \"{table}\";")
                row_count = conn.execute(row_count_query).scalar()
                # 获取字段个数
                column_count_query = text(f"""
                    SELECT COUNT(*) FROM information_schema.columns 
                    WHERE table_schema = 'public' AND table_name = '{table}';
                """)
                column_count = conn.execute(column_count_query).scalar()
                tables_info.append((table, row_count, column_count))
        return tables_info

    def get_table_ddl(self, table_name: str) -> str:
        """
        获取指定表的建表脚本 (DDL)
        :param table_name: 表名
        :return: 建表脚本字符串
        """
        with self.engine.connect() as conn:
            ddl_query = text(f"""
                SELECT pg_get_tabledef(oid) 
                FROM pg_class 
                WHERE relname = '{table_name}';
            """)
            result = conn.execute(ddl_query).fetchone()
            if result and result[0]:
                return result[0]
            else:
                raise ValueError(f"无法获取表 {table_name} 的建表脚本")


    def get_table_fields(self, table_name):
        """
        获取指定表的字段信息
        :param table_name: 表名称
        :return: 字段信息列表
        """
        query = text(f"""
        SELECT
            a.attname AS FieldName,
            col_description(a.attrelid, a.attnum) AS Comment,
            pg_catalog.format_type(a.atttypid, a.atttypmod) AS DataType
        FROM
            pg_catalog.pg_attribute a
        WHERE
            a.attnum > 0
            AND NOT a.attisdropped
            AND a.attrelid = (
                SELECT c.oid
                FROM pg_catalog.pg_class c
                JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
                WHERE c.relname = :table_name
                  AND n.nspname = 'public'
            );
        """)

        with self.engine.connect() as connection:
            result = connection.execute(query, {"table_name": table_name}).fetchall()
            return [{"FieldName": row[0], "Comment": row[1], "DataType": row[2],"IS_GEN": False} for row in result]

    def close(self):
        """
        关闭 SQLAlchemy 引擎
        """
        self.engine.dispose()